from PascalDataSet import PascalLabelImageDataSet, SingleImageDataSet, PascalDataSetGenerator
from torch.utils.data import DataLoader
import torch
from torchvision import transforms
from torchvision.transforms import Compose
import pickle
import os


def save_label(label, path: str):
    with open(path, 'wb') as f:
        pickle.dump(label, f)


if __name__ == "__main__":
    original_base = "/home/xiaoxiang/DataSet/VOC2012/"
    image_dir = os.path.join(original_base, "JPEGImages")
    label_dir = os.path.join(original_base, "SegmentationClass")

    target_base = "/home/xiaoxiang/DataSet/VOC2012_normalized/"
    target_image_dir = os.path.join(target_base, "JPEGImages")
    target_label_dir = os.path.join(target_base, "SegmentationClass")

    txt_dir = os.path.join(target_base, "Segmentation")

    generator = PascalDataSetGenerator("./color_object_dict.pickle")

    dataset = generator.create_dataset_by_txt("/home/xiaoxiang/DataSet/VOC2012/ImageSets/Segmentation/val.txt",
                                    image_dir, label_dir, crop_size=(320, 480))

    loader = DataLoader(dataset, batch_size=1, num_workers=8, shuffle=True)

    out_transformer = Compose([transforms.ToPILImage()])

    with open(os.path.join(txt_dir, "train.txt"), 'w') as f:
        for index, item in enumerate(loader):

            index = index+2000
            image = item[0]
            label = item[1]

            image = out_transformer(image[0])
            image.save(os.path.join(target_image_dir, str(index)+".jpg"))
            save_label(label, os.path.join(target_label_dir, str(index)+".pickle"))
            f.write(str(index)+"\n")

            print("finish: "+str(index)+"/"+str(dataset.__len__()))

    print("OK ")